import numpy as np

from src.impl.tf import get_tf_output
from src.impl.pytorch import get_torch_output
from src.impl.mnn import get_mnn_output

import tensorflow as tf


def get_output(tensor, operator, variable=None):
    if operator == 'avg_pool':
        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \
               get_mnn_output(tensor, operator).read(), None
    elif operator == 'max_pool':
        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \
               get_mnn_output(tensor, operator).read(), None
    elif operator == 'bias_add':
        if variable is None:
            variable = tf.random.normal([np.asarray(tensor).shape[-1]], stddev=0.01).numpy()
            return get_tf_output(tensor, operator, variable).numpy(), \
                get_torch_output(tensor, operator, variable).numpy(), \
                get_mnn_output(tensor, operator, variable).read(), variable
        else:
            return get_tf_output(tensor, operator, variable).numpy(), \
                get_torch_output(tensor, operator, variable).numpy(), \
                get_mnn_output(tensor, operator, variable).read(), None
    elif operator == 'conv2d':
        if variable is None:
            variable = tf.random.normal([3, 3, tensor.shape[-3], 32], stddev=0.01).numpy()
            return get_tf_output(tensor, operator, variable).numpy(), \
                get_torch_output(tensor, operator, variable).detach().numpy(), \
                get_mnn_output(tensor, operator, variable).read(), variable
        else:
            return get_tf_output(tensor, operator, variable).numpy(), \
                get_torch_output(tensor, operator, variable).detach().numpy(), \
                get_mnn_output(tensor, operator, variable).read(), None
    elif operator == 'softmax':
        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \
               get_mnn_output(tensor, operator).read(), None
    elif operator == 'batch_normalization':
        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \
               get_mnn_output(tensor, operator).read(), None
    elif operator == 'relu':
        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \
               get_mnn_output(tensor, operator).read(), None
    elif operator == 'reduce_mean':
        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \
               get_mnn_output(tensor, operator).read(), None
    elif operator == 'reduce_max':
        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \
               get_mnn_output(tensor, operator).read(), None
    elif operator == 'sigmoid':
        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \
               get_mnn_output(tensor, operator).read(), None
    elif operator == 'tanh':
        return get_tf_output(tensor, operator).numpy(), get_torch_output(tensor, operator).numpy(), \
               get_mnn_output(tensor, operator).read(), None
    elif operator == 'dense':
        if variable is None:
            variable = tf.random.normal([16, 10], stddev=0.01).numpy()
            return get_tf_output(tensor, operator, variable).numpy(), \
                   get_torch_output(tensor, operator, variable).numpy(), \
                   get_mnn_output(tensor, operator, variable).read(), variable
        else:
            return get_tf_output(tensor, operator, variable).numpy(), \
                   get_torch_output(tensor, operator, variable).numpy(), \
                   get_mnn_output(tensor, operator, variable).read(), None


def get_output_float16(tensor, operator, variable):
    tensor_16 = tensor.astype(np.float16)
    return get_output(tensor_16, operator, variable)


def get_output_float32(tensor, operator):
    tensor_32 = tensor.astype(np.float32)
    return get_output(tensor_32, operator)


def get_output_float32_with_variable(tensor, operator, variable):
    tensor_32 = tensor.astype(np.float32)
    return get_output(tensor_32, operator, variable)